﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.IO;

namespace Kokomo.Mirror.Ia32
{
	public class Ia32Reader
	{
		public Ia32Reader(Stream stream, long streamBaseAddress)
		{
			if (stream == null) throw new ArgumentNullException("stream");

			this.Stream = stream;
			this.StreamBaseAddress = streamBaseAddress;
			this.Reader = new BinaryReader(stream);
		}

		public long StreamBaseAddress { get; private set; }
		public Stream Stream { get; private set; }
		protected BinaryReader Reader { get; private set; }
		public Ia32Instruction CurrentInstruction { get; private set; }

		public bool MoveNext()
		{
			long startStreamOffset = this.Stream.Position;
			long startOffset = startStreamOffset + this.StreamBaseAddress;

			Prefixes prefixes = 0;
			Register segmentOverrideRegister = null;

			int opCodeSize = 1;
			int opCodeValue = 0;
			while (true)
			{
				int n = this.Stream.ReadByte();
				if (n == -1) return false;

				opCodeValue |= n;
				if (n == OpCodes.EscapeCode)
				{
					opCodeSize = 2;
					opCodeValue <<= 8;
					continue;
				}

				if (OpCodes.IsPrefix(n))
				{
					prefixes |= OpCodes.GetPrefixCode(n);
					opCodeValue = 0;

					switch (prefixes)
					{
						case Prefixes.CsSegmentOverride:
							segmentOverrideRegister = Registers.CS;
							break;
						case Prefixes.SsSegmentOverride:
							segmentOverrideRegister = Registers.SS;
							break;
						case Prefixes.DsSegmentOverride:
							segmentOverrideRegister = Registers.DS;
							break;
						case Prefixes.EsSegmentOverride:
							segmentOverrideRegister = Registers.ES;
							break;
						case Prefixes.FsSegmentOverride:
							segmentOverrideRegister = Registers.FS;
							break;
						case Prefixes.GsSegmentOverride:
							segmentOverrideRegister = Registers.GS;
							break;
					}

					continue;
				}

				bool hasModRM = false;
				ModRM modrm = new ModRM();
				OpCode opCode;
				if (OpCodes.IsExtendedOpCode(opCodeValue))
				{
					modrm = new ModRM(this.Reader.ReadByte());
					hasModRM = true;

					opCode = OpCodes.GetExtendedOpCode(opCodeValue, modrm, prefixes);
				}
				else
				{
					opCode = OpCodes.GetMatchingOpCode(opCodeValue, opCodeSize);
				}

				if (opCode.IsUnknown) return false;

				ScaleIndexBase sib = new ScaleIndexBase();
				if (opCode.Form.HasFlag(OpCodeForm.ModRM) && !hasModRM)
				{
					modrm = new ModRM((byte)this.Stream.ReadByte());
					hasModRM = true;
				}
				int displacementSizeBytes = 0;
				if (hasModRM)
				{
					Mod mod = modrm.Mod;
					if (mod == Mod.disp8) displacementSizeBytes = 1;
					else if (mod == Mod.disp32 || modrm.AddressingForm32 == AddressingForm32.disp32)
						displacementSizeBytes = 4;

					if (modrm.IndicatesSib)
					{
						sib = new ScaleIndexBase((byte)this.Stream.ReadByte());
						if (
							(sib.Base == Base.disp)
							&& (displacementSizeBytes == 0)
							)
						{
							displacementSizeBytes = 4;
						}
					}
				}

				int displacement = 0;
				switch (displacementSizeBytes)
				{
					case 0:
						break;
					case 1:
						displacement = this.Reader.ReadSByte();
						break;
					case 4:
						displacement = this.Reader.ReadInt32();
						break;
					default:
						throw new InvalidOperationException("Unknown displacement");
				}

				Operand[] operands = new Operand[opCode.OperandTypes.Length];
				for (int i = 0; i < operands.Length; i++)
				{
					OperandType operandType = opCode.OperandTypes[i];

					OperandType operandEncoding = operandType & OperandType.EncodingMask;
					OperandType operandDataType = operandType & OperandType.DataTypeMask;
					OperandType operandSemantic = operandType & OperandType.SemanticMask;

					int sizeBytes = this.DetermineOperandSize(operandDataType, prefixes);

					Operand operand;
					switch (operandEncoding)
					{
						case OperandType.OpCodeReg:
							{
								int registerNumber = n & 0x07;
								RegisterType registerType = (RegisterType)((int)operandSemantic >> (int)OperandType.RegisterTypeShift);
								Register register = Register.Get(registerNumber, sizeBytes * 8, registerType);
								if (register == null) throw new InvalidOperationException("Unrecognized register");
								operand = register;
							}
							break;
						case OperandType.ModRM:
							if (modrm.Mod == Mod.reg)
							{
								RegisterType registerType;
								switch (operandSemantic)
								{
									case OperandType.GeneralPurposeRegister:
									case OperandType.GprOrAddress:
										registerType = RegisterType.GeneralPurpose;
										break;
									case OperandType.Mmx:
									case OperandType.MmxOrAddress:
										registerType = RegisterType.Mmx;
										break;
									case OperandType.Xmm:
									case OperandType.XmmOrAddress:
										registerType = RegisterType.Simd;
										break;
									case OperandType.FpuRegister:
									case OperandType.FpuOrAddress:
										registerType = RegisterType.Fpu;
										break;
									case OperandType.ControlRegister:
										registerType = RegisterType.Control;
										break;
									default:
										throw new InvalidOperationException("Unrecognized operand");
								}

								Register register = Register.Get(modrm.RM, sizeBytes * 8, registerType);
								if (register == null) throw new InvalidOperationException("Unrecognized operand");

								operand = register;
							}
							else
							{
								Address32Reference address = new Address32Reference(
									modrm.AddressingForm32,
									sizeBytes,
									segmentOverrideRegister,
									displacement,
									sib.ScaleFactor,
									sib.Index,
									sib.Base
									);
								operand = address;
							}
							break;
						case OperandType.Reg:
							{
								int registerNumber = modrm.Reg;
								Register register = Register.Get(registerNumber, sizeBytes * 8, (RegisterType)((int)operandSemantic >> (int)OperandType.RegisterTypeShift));
								if (register == null) throw new InvalidOperationException("Unrecognized register");
								operand = register;
							}
							break;
						case OperandType.FixedRegister:
							{
								int registerCode = ((int)operandType & (int)OperandType.FixedRegisterMask) >> (int)OperandType.FixedRegisterShift;
								RegisterType registerType = (RegisterType)((int)operandSemantic >> (int)OperandType.RegisterTypeShift);
								Register register = Register.Get(registerCode, sizeBytes * 8, registerType);
								if (register == null) throw new InvalidOperationException("Unrecognized register");
								operand = register;
							}
							break;
						case OperandType.ImplicitPointerDI:
							operand = new Address32Reference(AddressingForm32.edi, sizeBytes);
							break;
						case OperandType.ImplicitPointerSI:
							operand = new Address32Reference(AddressingForm32.esi, sizeBytes);
							break;
						case OperandType.Immediate:
							{
								switch (operandSemantic)
								{
									case OperandType.JumpRelativeOffset:
										int relativeOffset;
										switch (sizeBytes)
										{
											case 1:
												relativeOffset = this.Reader.ReadSByte();
												break;
											case 2:
												relativeOffset = this.Reader.ReadInt16();
												break;
											case 4:
												relativeOffset = this.Reader.ReadInt32();
												break;
											default:
												throw new InvalidOperationException("Unrecognized immediate size");
										}

										operand = new JumpTarget(startOffset, Stream.Position + this.StreamBaseAddress + relativeOffset);
										break;
									case OperandType.JumpAbsoluteOffset:
									case OperandType.Address:
										uint address;
										int addressSize = this.DetermineAddressSize(prefixes);
										switch (addressSize)
										{
											case 1:
												address = this.Reader.ReadByte();
												break;
											case 2:
												address = this.Reader.ReadUInt16();
												break;
											case 4:
												address = this.Reader.ReadUInt32();
												break;
											default:
												throw new InvalidOperationException("Unrecognized immediate size");
										}

										if (operandSemantic == OperandType.JumpAbsoluteOffset)
											operand = new JumpTarget(startOffset, address);
										else
											operand = new Address32Reference(AddressingForm32.disp32, sizeBytes, segmentOverrideRegister, (int)address);

										break;
									default:
										uint immediateValue;
										switch (sizeBytes)
										{
											case 1:
												immediateValue = this.Reader.ReadByte();
												break;
											case 2:
												immediateValue = this.Reader.ReadUInt16();
												break;
											case 4:
												immediateValue = this.Reader.ReadUInt32();
												break;
											default:
												throw new InvalidOperationException("Unrecognized immediate size");
										}

										Immediate immediate = new Immediate(immediateValue);
										operand = immediate;
										break;
								}
							}
							break;
						default:
							throw new InvalidOperationException("Unrecognized operand type");
							break;
					}

					operands[i] = operand;
				}

				long endStreamOffset = this.Stream.Position;
				int instructionLength = (int)(endStreamOffset - startStreamOffset);
				byte[] bytes = new byte[instructionLength];
				this.Stream.Position = startStreamOffset;
				this.Stream.Read(bytes, 0, instructionLength);
				Ia32Instruction instruction = new Ia32Instruction(
					bytes,
					startOffset,
					instructionLength,
					prefixes,
					opCode,
					operands
					);

				this.CurrentInstruction = instruction;

				return true;
			}
		}

		private int DetermineAddressSize(
			Prefixes prefixes
			)
		{
			int size = prefixes.HasFlag(Prefixes.AddressSizeOverride)
				? 2
				: 4
				;
			return size;
		}

		private int DetermineOperandSize(
			OperandType size,
			Prefixes prefixes
			)
		{
			int sizeBytes;

			bool hasOperandSizeOverridePrefix = prefixes.HasFlag(Prefixes.OperandSizeOverride);
			bool hasAddressSizeOverridePrefix = prefixes.HasFlag(Prefixes.AddressSizeOverride);

			switch (size)
			{
				case OperandType.Byte:
					sizeBytes = 1;
					break;
				case OperandType.Word:
					sizeBytes = 2;
					break;
				case OperandType.Dword:
					sizeBytes = 4;
					break;
				case OperandType.Qword:
				case OperandType.DDword:
					sizeBytes = 8;
					break;
				case OperandType.DQWord:
					sizeBytes = 16;
					break;
				case OperandType.PackedSingle:
					sizeBytes = 16;
					break;
				case OperandType.OperandByteOrWord:
					sizeBytes = (hasOperandSizeOverridePrefix ? 1 : 2);
					break;
				case OperandType.OperandWordOrDword:
				case OperandType.OperandWordDwordQword:
					sizeBytes = (hasOperandSizeOverridePrefix ? 2 : 4);
					break;
				case OperandType.OperandDwordOrFword:
					sizeBytes = 4;
					break;
				case OperandType.AddressWordOrDword:
					sizeBytes = (hasAddressSizeOverridePrefix ? 2 : 4);
					break;
				case OperandType.DoubleFloatingPoint:
					sizeBytes = 8;
					break;
				case OperandType.PseudoDescriptor:
					sizeBytes = 6;
					break;
				default:
					throw new ArgumentException("Invalid operand size");
			}

			return sizeBytes;
		}
	}
}
